package app.dbunit;

import core.framework.util.StringUtils;
import org.dbunit.database.IDatabaseConnection;
import org.dbunit.dataset.CompositeDataSet;
import org.dbunit.dataset.DataSetException;
import org.dbunit.dataset.FilteredDataSet;
import org.dbunit.dataset.IDataSet;
import org.dbunit.dataset.ITable;
import org.dbunit.dataset.filter.ExcludeTableFilter;
import org.dbunit.dataset.xml.FlatXmlDataSetBuilder;
import org.dbunit.ext.hsqldb.HsqldbConnection;
import org.dbunit.ext.mssql.MsSqlConnection;
import org.dbunit.ext.mysql.MySqlConnection;
import org.dbunit.operation.DatabaseOperation;
import org.dbunit.operation.DatabaseOperationExt;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.sql.Connection;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author Bribin
 */
public class DBUnitTool {

    private final Logger logger = LoggerFactory.getLogger(DBUnitTool.class);

    public static final String DBTYPE_HSQL = "HSQL";
    public static final String DBTYPE_MSSQL = "MSSQL";
    public static final String DBTYPE_MYSQL = "MYSQL";

    private static final Map<String, List<String>> DB_TYPES = new HashMap<String, List<String>>();

    static {
        DB_TYPES.put(DBTYPE_HSQL, Arrays.asList("hsql"));
        DB_TYPES.put(DBTYPE_MSSQL, Arrays.asList("mssql", "Microsoft SQL Server"));
        DB_TYPES.put(DBTYPE_MYSQL, Arrays.asList("mysql"));
    }

    private String dbType;
    private IDatabaseConnection connection;

    public DBUnitTool(Connection conn) throws Exception {
        judgeDBType(conn);

        if (DBTYPE_HSQL.equals(dbType))
            connection = new HsqldbConnection(conn, conn.getSchema());
        else if (DBTYPE_MSSQL.equals(dbType))
            connection = new MsSqlConnection(conn);
        else if (DBTYPE_MYSQL.equals(dbType))
            connection = new MySqlConnection(conn, null);
    }

    private void judgeDBType(Connection conn) throws SQLException {
        String productName = conn.getMetaData().getDatabaseProductName().toLowerCase();
        for (Map.Entry<String, List<String>> entry : DB_TYPES.entrySet()) {
            for (String type : entry.getValue()) {
                if (productName.indexOf(type.toLowerCase()) >= 0) {
                    dbType = entry.getKey();
                    break;
                }
            }
        }

        logger.debug("detect database type: {}", dbType);

        // TODO(Bribin) to support other database type
        if (dbType == null)
            throw new IllegalArgumentException("This database product [" + productName + "] is not supported yet !");
    }

    public IDatabaseConnection getConnection() {
        return connection;
    }

    private DatabaseOperation getInsertOperation() {
        if (DBTYPE_MSSQL.equals(dbType))
            return DatabaseOperationExt.IDENTITY_INSERT;
        return DatabaseOperationExt.INSERT;
    }

    private DatabaseOperation getCleanInsertOperation() {
        if (DBTYPE_MSSQL.equals(dbType))
            return DatabaseOperationExt.IDENTITY_CLEAN_INSERT;
        return DatabaseOperationExt.CLEAN_INSERT;
    }

    private DatabaseOperation getRefreshOperation() {
        if (DBTYPE_MSSQL.equals(dbType))
            return DatabaseOperationExt.IDENTITY_REFRESH;
        return DatabaseOperationExt.REFRESH;
    }

    private DatabaseOperation getTruncateOperation() {
        return DatabaseOperationExt.TRUNCATE_TABLE;
    }

    public void cleanInertDataSet(String... datafiles) throws Exception {
        logger.debug("create datas with files: " + Arrays.toString(datafiles));
        databaseOperation(getCleanInsertOperation(), datafiles);
    }

    public void createDataSet(String... datafiles) throws Exception {
        logger.debug("create datas with files: " + Arrays.toString(datafiles));
        databaseOperation(getInsertOperation(), datafiles);
    }

    public void updateDataSet(String... datafiles) throws Exception {
        logger.debug("update datas with files: " + Arrays.toString(datafiles));
        databaseOperation(getRefreshOperation(), datafiles);
    }

    public void deleteDataSet(String... datafiles) throws Exception {
        logger.debug("delete datas with files: " + Arrays.toString(datafiles));
        databaseOperation(DatabaseOperation.DELETE, datafiles);
    }

    private void databaseOperation(DatabaseOperation operation, String[] datafiles) throws Exception {
        if (datafiles == null || datafiles.length == 0)
            return;
        operation.execute(connection, buildDataSet(datafiles));
    }

    private IDataSet buildDataSet(String[] datafiles) throws IOException, DataSetException {
        List<IDataSet> dataSets = new ArrayList<IDataSet>();
        FlatXmlDataSetBuilder builder = new FlatXmlDataSetBuilder();
        builder.setColumnSensing(true);
        for (String datafile : datafiles) {
            if (datafile != null && !datafile.isEmpty() && getClass().getClassLoader().getResource(datafile) != null) {
                logger.debug("to load datafile: " + datafile);
                dataSets.add(builder.build(getClass().getClassLoader().getResourceAsStream(datafile)));
            }
        }
        return new CompositeDataSet(dataSets.toArray(new IDataSet[dataSets.size()]), false);
    }

    public void truncateAllTables() throws Exception {
        IDataSet dataSet = connection.createDataSet();

        if (DBTYPE_MSSQL.equals(dbType)) {
            ExcludeTableFilter filter = new ExcludeTableFilter(new String[]{"*trace_xe_action_map", "*trace_xe_event_map"});
            dataSet = new FilteredDataSet(filter, dataSet);
        }

        // TODO(Bribin) due to foreign key constraints, the table names should be sorted
        logger.info("to truncate all tables: {}", Arrays.toString(dataSet.getTableNames()));

        getTruncateOperation().execute(connection, dataSet);
    }

    public ITable query(String sql) throws Exception {
        logger.debug("query SQL: {}", sql);
        return connection.createQueryTable("query", sql);
    }

    public void execute(String sql) throws Exception {
        logger.debug("execute SQL: {}", sql);
        if (StringUtils.hasText(sql)) {
            try (Statement stmt = connection.getConnection().createStatement()) {
                stmt.execute(sql);
            } catch (SQLException e) {
                logger.error(e.getMessage());
            }
        }
    }

}
